#-*- coding:utf-8 -*-
import numpy as np
from scipy.linalg import solve
import matplotlib.pyplot as plt

'''
读取数据
数据格式:
网格数	传热系数	横截面积	管路长度	西端温度	东端温度	热源
5		0.5 		1 			0.02 		100 		200			1e6
'''
def read_datas(datafile):
    f=open(datafile)
    datafiles=[]
    for line in f:
        datafiles.append(line.split())
    f.close()
    return datafiles

#写入结果
def writeDatas(fileName,params):
	file_object = open(fileName, 'w')
	file_object.write('计算结果:\n')
	for y in params['y']:
		file_object.write('{0:.2f}\n'.format(y))
	file_object.close()

#输出对齐
def myAlign(string, length=0):  
    if length == 0:  
        return string  
    slen = len(string)  
    re = string  
    if isinstance(string, str):  
        placeholder = ' '  
    else:  
        placeholder = u'　'  
    while slen < length:  
        re += placeholder  
        slen += 1  
    return re 

#初始化数据输出
def printDatas(datas):
	for i in range(2):
		for data in datas[i]:
			print(myAlign(data,10)),
		print('\n'),

#结果输出
def printResult(params):
	print(u'AX = B中的A:')
	print(params['a'])
	print(u'AX = B中的B:')
	print(params['b'])
	print(u'所求解X:')
	print(params['y'])

#图形绘制
def creat_plot(params):
	#x = np.linspace(0, params['L'] , params['N'])
	x = [0.002,0.006,0.01,0.014,0.018]
	x = np.arange(0, 0.02, params['L'] / params['N'])
	y = params['y']
	#y0 = np.array([])
	#y0.resize(params['N'] + 2)
	#for i in range(y.size):
	#	y0[i+1] = y[i]
	#y0[0] = 100
	#y0[params['N']+1] = 500
	#y = y0
	
	z1 = np.polyfit(x, y, 2)
	# 生成的多項式對象
	p1 = np.poly1d(z1)
 
	# 第2個擬合，自由度為6
	z2 = np.polyfit(x, y, 6)
	# 生成的多項式對象
	p2 = np.poly1d(z2)
	
	plt.plot(x, y, 'r*')
	plt.plot(x,p1(x),'b-')
	plt.ylabel("temperature($C$)")
	plt.xlabel("Length($m$)")
	plt.ylim(100,260)
	plt.xlim(0,0.02)
	plt.grid()
	#plt.title("PyPlot First Example")
	#plt.ylim(-1.2,1.2)
	#plt.legend()
	plt.show()

def cacul():
	datas = read_datas("datasq.txt")
	printDatas(datas)

	params = {}
	#数据格式转换
	params['N'] = int(datas[1][0])
	params['k'] = float(datas[1][1])
	params['A'] = float(datas[1][2])
	params['L'] = float(datas[1][3])
	params['TA'] = float(datas[1][4])
	params['TB'] = float(datas[1][5])
	params['Q'] = float(datas[1][6])
	#Δx计算
	dx = params['L'] / params['N']

	#将相同计算作为一个系数
	factor =  params['k'] * params['A'] / dx
	qADx = params['Q'] * params['A'] * dx

	#Ax = B
	#初始化A,B
	a = np.array([])
	b = np.array([])
	a.resize(params['N'],params['N'])
	b.resize(params['N'])

	#赋值给B
	b[0] = 2 * factor * params['TA'] + qADx
	for i in range(1,params['N']-1):
		b[i] = qADx
	b[params['N']-1] = 2 * factor * params['TB'] + qADx

	#第一个点和最后一个点计算
	a[params['N']-1][params['N']-1] =a[0][0] = 3 * factor
	a[params['N']-1][params['N']-2] = a[0][1] = - factor

	#中间点计算
	for i in range(1,params['N']-1):
		a[i][i-1] =a[i][i+1] =  - factor
		a[i][i] = 2 * factor
	y = solve(a, b)
	params['a'] = a
	params['b'] = b
	params['y'] = y
	return params

if __name__ == '__main__': 
	params = cacul()
	printResult(params)
	writeDatas('resultq.txt',params)
	creat_plot(params)

